#########################################################################
## https://greenleaflab.github.io/chromVAR/articles/Articles/Applications.html#differential-accessibility-and-variability

library(chromVAR)
library(chromVARmotifs)
library(motifmatchr)
library(SummarizedExperiment)
library(BSgenome.Hsapiens.UCSC.hg19)
library(BSgenome.Hsapiens.UCSC.hg38)
library(BiocParallel)
library(Matrix)
library(pheatmap)
library(patchwork)
library(parallel)
library(dplyr)
library(optparse)

##########################################################################################

option_list <- list(
    make_option(c("--peakfile"), type = "character"),
    make_option(c("--cell_qc_file"), type = "character"),
    make_option(c("--bamfiles"), type = "character"),
    make_option(c("--cluster"), type = "character"),
    make_option(c("--celltype_name"), type = "character"),
    make_option(c("--genome_type"), type = "character"),
    make_option(c("--data_base"), type = "character"),
    make_option(c("--cpu"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    setwd("~/20231121_singleMuti/input")
    peakfile <- "cluster0_peak.bed"
    bamfiles <- c("cluster0_merge_sorted.reviewRG.bam")
    celltype_name <- "Myoid"
    genome_type <- "hg38"
    out_path <- "~/20231121_singleMuti/results"
    data_base <- "jaspar"
    cpu <- 10
    cluster <- "cluster0"
    cell_qc_file <- "~/20231121_singleMuti/results/cluster_all_result/cell_fragment/testis_merge_all_qc_barcode_new.txt"
}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

cluster <- opt$cluster
peakfile <- opt$peakfile
cell_qc_file <- opt$cell_qc_file
bamfiles <- opt$bamfiles
celltype_name <- opt$celltype_name
genome_type <- opt$genome_type
data_base <- opt$data_base
cpu <- opt$cpu
out_path <- opt$out_path

#########################################################################
## 基因组版本
if(genome_type == "hg38"){
    genome_use <- BSgenome.Hsapiens.UCSC.hg38
}else{
    genome_use <- BSgenome.Hsapiens.UCSC.hg19
}

## TF版本
if(data_base == "jaspar"){
    motifs <- getJasparMotifs(species = "Homo sapiens")
}else if(data_base == "cisbp"){
    data("human_pwms_v2")
    motifs <- human_pwms_v2
}

## 质控后的细胞
cell_qc <- read.csv(cell_qc_file)

# 开起多核处理
register(MulticoreParam(cpu))
dir.create(out_path , recursive = T)

#########################################################################
### how to read in counts
peaks <- getPeaks(peakfile, sort_peaks = TRUE)
fragment_counts <- getCounts(bamfiles, peaks, 
                              paired = TRUE, 
                              by_rg = TRUE, 
                              format = "bam", 
                              colData = DataFrame(celltype = celltype_name))

## 提取质控通过的细胞
fragment_counts@colData <- fragment_counts@colData[rownames(fragment_counts@colData) %in% cell_qc$V1,]

#########################################################################
## https://mp.weixin.qq.com/s/_qqTvQH085OTmFOZG0Maxg
#### 1、Getting GC content of peaks
## GC内容将用于确定背景peaks。函数addGCBias将在rowData中添加名为“bias”的新列，并返回一个更新后的SummarizedExperiment 对象。
## 该函数需要一个基因组序列的输入，可以由BSgenome、FaFile或DNAStringSet等对象来提供。
counts_filtered <- addGCBias(fragment_counts , genome = genome_use)

#### 2、质控
## 如果处理单个细胞数据，建议过滤出 reads 不足或peaks 中 reads 比例较低的样本，因为这些样本可能代表空或死细胞。
## 两个参数用于过滤—min_in_peak和min_depth。如果没有提供(如上所述)，则根据数据中的中位数来估计这些 cutoffs 。 
## min_in_peaks 被设置为peak中fragment比例中位数的0.5倍。
## min_depth 被设置为最大的500或中位数库大小的10%。
counts_filtered <- filterSamples(counts_filtered , shiny = FALSE)

## 对于 bulk 和单细胞数据， peaks 都应该根据至少有一定数量的 fragments 进行筛选。
## 至少，每个peaks应该在所有样本中至少有一个fragments片段(由于使用由其他数据定义的peaks，可能出现reads为零的peaks)。否则，下游功能将无法工作。
## 如果non_overlap参数被设置为TRUE(这是默认值)，那么filterPeaks也会将peaks设置为不重叠的peaks(对于重叠的peaks，会保留计数更高的peaks)。
counts_filtered <- filterPeaks(counts_filtered)

#### 3、Get motifs and what peaks contain motifs
# data("human_pwms_v2") #filtered collection of human motifs from cisBP database
# data("mouse_pwms_v2") #filtered collection of mouse motifs from cisBP database
# data("homer_pwms") #motifs from HOMER
# data("encode_pwms") #motifs from ENCODE

## https://bioconductor.org/packages/release/bioc/vignettes/motifmatchr/inst/doc/motifmatchr.html
## 三种输出解释
## (Default, with out = "matches") Boolean matrix indicating which ranges/sequences contain which motifs, 
## stored as "matches" in assays slot of SummarizedExperiment object
## (out = "scores") Same as (1) plus two additional assays -- a matrix with the score of the high motif score within each range/sequence 
## (score only reported if match present) and a matrix with the number of motif matches.
## (out = "positions") A GenomicRangesList with the ranges of all matches within the input ranges/sequences.
## p.cutoff，用于确定模体调用应该有多严格。默认值为0.00005，通常能够给出人类模体的合理匹配数量。
## 除了返回模体匹配之外，该函数还可以返回额外的矩阵（存储为 assays），其中包含每个峰值的模体匹配数量和每个峰值的最大模体分数。
## 对于这些额外的信息，请使用 out = scores。要返回模体匹配的实际位置，
## 请使用 out = positions。 out = matches 或 out = scores 的输出都可以传递给 computeDeviations 函数。

motif_ix <- matchMotifs(motifs, counts_filtered , genome = genome_use)
motif_ix_scores <- matchMotifs(motifs, counts_filtered , genome = genome_use , out = "scores")
motif_ix_positions <- matchMotifs(motifs, counts_filtered , genome = genome_use , out = "positions")


#### 4、computing deviations
## 背景peaks对偏差值进行归一化
dev <- computeDeviations(object = counts_filtered , annotations = motif_ix)

#########################################################################
#### 提取motif和cell的z-score
## The function `computeDeviations` returns a SummarizedExperiment with two "assays". 
## The first matrix (accessible via `deviations(dev)` or `assays(dev)$deviations`) will give the bias corrected "deviation" in 
## accessibility for each set of peaks (rows) for each cell or sample (columns). 
## This metric represent how accessible the set of peaks is relative to the expectation based on equal chromatin accessibility profiles across cells/samples, 
## normalized by a set of background peak sets matched for GC and average accessability.  
## The second matrix (`deviationScores(dev)` or `assays(deviations)$z`) gives the deviation Z-score, 
## which takes into account how likely such a score would occur if randomly sampling sets of beaks with similar GC content and average accessibility.  

motif_cell_score <- deviationScores(dev)
out_name <- paste0(out_path , "/motif_cell.z-score.tsv")
write.table(motif_cell_score , out_name , row.names = T , quote = F , sep = "\t")

#########################################################################
#### 提取motif和peak的对应关系
## The package motifmatchr can be used to find motifs within peaks.  
## The `matchMotifs` method takes as inputs the motif lists described above, or your own list of motifs, 
## and returns a (Ranged)SummarizedExperiment with a matrix indicating what peaks (rows) contain what motifs (columns). 
## 提取与每个peak相匹配的motif
peak_dat <- data.frame(motif_ix@rowRanges)
result_motif_peak <- Reduce(function(x,y)bind_rows(x,y),mclapply(1:nrow(motif_ix),function(i){
    print(i)
    motifs <- names(which(assay(motif_ix)[i,]))
    if(length(motifs) > 0){
        tmp_peak_dat <- cbind(peak_dat[i,] , motifs)
    }else{
        tmp_peak_dat <- data.frame()
    }
    tmp_peak_dat
},mc.cores=cpu))

out_name <- paste0(out_path , "/motif_peak.tsv")
result_motif_peak$start <- result_motif_peak$start-1
result_motif_peak$width <- result_motif_peak$end - result_motif_peak$start
write.table(result_motif_peak , out_name , row.names = F , quote = F , sep = "\t")

#########################################################################
#### 保存为Rdata
out_name <- paste0(out_path , "/" , cluster , ".chromvar.rda")
save(dev, motif_ix , motif_ix_scores , motif_ix_positions , file = out_name)

#########################################################################
if(1!=1){
    #### 以下不需要运行，用于多个不同类型的细胞间的比较
    #########################################################################
    #################################
    ## Variability  计算感兴趣的cell或样本之间每个motif或其他注释的变异性
    ## computeVariability函数返回一个包含变异性data.frame( 计算了所有细胞/样本一组 peaks的z分数的标准差), 
    ## bootstrap 置信区间的变化(由重采样细胞/样品),和 变异性大于原假设1的p值
    variability <- computeVariability(dev)
    out_name <- paste0(out_path , "/variability.pdf")
    pdf(out_name)
    plotVariability(variability, use_plotly = FALSE)
    dev.off()
    #################################
    ## 聚类分析
    ## 也可以用偏差校正偏差（bias corrected deviations）对样本聚类，
    ## 函数getSampleCorrelation首先剔除高度相关注释和低可变性的注释，然后计算细胞间的相关性
    ## 相关性计算
    sample_cor <- getSampleCorrelation(dev)

    out_name <- paste0(out_path , "/pheatmap_SampleCorrelation.pdf")
    pdf(out_name)
    pheatmap(
        as.dist(sample_cor), 
        annotation_row = colData(dev), 
        clustering_distance_rows = as.dist(1-sample_cor), 
        clustering_distance_cols = as.dist(1-sample_cor))
    dev.off()

    #################################
    ## 细胞或样本间的相似性 Visualizing Deviations
    ## 对于细胞的可视化，可以使用TSNE将偏差值投射到二维中。ChromVAR提供了一个方便的函数deviationsTsne。
    ## perplexity根据细胞的实际数目调整，通常情况下，100个细胞时，perplexity在30-50都是可以的
    tsne_results <- deviationsTsne(dev, threshold = 1.5, perplexity = 10)

    tsne_plots <- plotDeviationsTsne(dev, tsne_results, annotation_name = "TEAD3", sample_column = "Cell_Type", shiny = FALSE)

    out_name <- paste0(out_path , "/TEAD3_DeviationsTsne.pdf")
    pdf(out_name , width = 9 , height = 5)
    tsne_plots[[1]] + tsne_plots[[2]]
    dev.off()


    #################################
    ## 可及性和变异性差异
    ## 差分偏差函数确定不同组之间给定注释的偏差校正偏差之间是否存在显着差异

    diff_acc <- differentialDeviations(dev, "Cell_Type")
    out_name <- paste0(out_path , "/differentialDeviations.tsv")
    write.table(diff_acc , out_name , row.names = F , quote = F , sep = "\t")

    diff_var <- differentialVariability(dev, "Cell_Type")
    out_name <- paste0(out_path , "/differentialVariability.tsv")
    write.table(diff_var , out_name , row.names = F , quote = F , sep = "\t")

    #################################
    ## motif/kmer
    ## We can also perform tsne for motif similarity rather than cell similarity, by specifying what = "annotations" to the deviationsTsne function.
    inv_tsne_results <- deviationsTsne(dev, threshold = 1.5, perplexity = 8, 
                                        what = "annotations", shiny = FALSE)

    p <- ggplot(inv_tsne_results, aes(x = Dim1, y = Dim2)) + geom_point() + 
      chromVAR_theme()

    out_name <- paste0(out_path , "/deviationsTsne.pdf")
    ggsave(out_name , p , width = 5 , height = 3 )

}